from django.test import TestCase
from database.chaos_mesh import get_stressor_type, get_service_for_pod, str2int_time, ChaosMeshConnection
from database import cm_conn

import pymysql
import json
from datetime import datetime, timedelta

class ChaosMeshTestCase(TestCase):
    def test_get_stressor_type(self):
        stressors = {"io": "high", "cpu": "medium"}
        self.assertEqual(get_stressor_type(stressors), "io")

    def test_get_service_for_pod(self):
        cmdb_id_list = ["service-pod1", "service-pod2"]
        self.assertEqual(get_service_for_pod(cmdb_id_list), ["service", "service"])

    def test_str2int_time(self):
        # duration_str[-1] in units
        self.assertEqual(str2int_time("5m"), 300)
        # duration_str[-1] not in units
        with self.assertRaises(ValueError):
            str2int_time("5x")
        # wrong format: <number><unit>
        with self.assertRaises(ValueError):
            str2int_time("am")

    def test_ChaosMeshConnection_init(self):
        right_config = {
            "host": "1.92.152.201",
            "user": "root",
            "password": "elastic",
            "port": 3306,
            "database": "chaos_mesh",
            "kubernetes_path": "/home/group04/containers/aiops-backend/kubeconfig/config",
        }
        error_config = {
            "host": "1.92.152.201",
            "user": "root",
            "password": "elasticelastic",
            "port": 3306,
            "database": "chaos_mesh",
            "kubernetes_path": "/home/group04/containers/aiops-backend/kubeconfig/config",
        }
        # 连接成功
        connection_right = ChaosMeshConnection(right_config)
        self.assertIsNotNone(connection_right.connection)
        # 连接失败
        with self.assertRaises(pymysql.Error):
            connection_error = ChaosMeshConnection(error_config)

    def test_ChaosMeshConnection_get_pods_for_label(self):
        namespace = "default"
        label = "productcatalogservice"
        
        self.assertEqual(cm_conn.get_pods_for_label(namespace, label), ["productcatalogservice-5dff754586-gjtdf"])

    def test_ChaosMeshConnection_get_culprit_cmdb_id(self):
        # if not "pods" in selector:
        cmdb_id_list = ['frontend-69c95fbb99-4tbkt']
        selectors = {'namespaces': ['default'], 'labelSelectors': {'app': 'frontend'}}
        self.assertEqual(cm_conn.get_culprit_cmdb_id(cmdb_id_list, selectors), ['frontend-69c95fbb99-4tbkt'])
        # if "pods" in selector:
        cmdb_id_list = ['pods_in_list']
        selectors = {'namespaces': ['default'], 'pods': {"default": ['pods_in_list', 'pods_not_in_list']}}
        self.assertEqual(cm_conn.get_culprit_cmdb_id(cmdb_id_list, selectors), ['pods_in_list'])
        
    def test_ChaosMeshConnection_find_duration_in_workflow(self):
        pass
    
    def test_ChaosMeshConnection_feature_extract(self):
        data = {}
        with open("./api/tests/choaos_mesh.json", "r") as f:
            data = json.load(f)
        result = (["recommendationservice"], ["recommendationservice-778fc56466-r6gkl"], "delay", 900)
        result3 = (['recommendationservice', 'recommendationservice'], ['recommendationservice-778fc56466-r6gkl', 'recommendationservice-778fc56466-r6gkl'], 'delay', 900)

        # 共 5 种情况，覆盖所有语句分支，如下：
        # "duration" in experiment.keys(), kind == "NetworkChaos", direction == "to"
        # "duration" in experiment.keys(), kind == "NetworkChaos", direction == "from"
        # "duration" in experiment.keys(), kind == "NetworkChaos", direction != "to" 和 "from"
        # "duration" in experiment.keys(), kind == "StressChaos"
        # "duration" in experiment.keys(), kind == "PodChaos"

        for key in data["feature_extract_data"].keys():
            cur_test_data = data["feature_extract_data"][key]
            cur_result = result
            if key == "test3":
                cur_result = result3
            self.assertEqual(cm_conn.feature_extract(cur_test_data["experiment"], cur_test_data["cmdb_id_list"], cur_test_data["kind"], cur_test_data["metadata"]), cur_result)

    def test_ChaosMeshConnection_group_by_object_id(self):
        start_time = datetime(2024, 5, 19, 8, 20, 0)
        end_time = datetime(2024, 5, 19, 8, 21, 0)
        result = {'929f3a12-bc6c-425a-a2ba-524ed6f2b959': {'object_id': '929f3a12-bc6c-425a-a2ba-524ed6f2b959', 'created_at': 1716106800, 'kind': 'PodChaos', 'cmdb_id': ['checkoutservice-6846846596-tff57']}}
        self.assertEqual(cm_conn.group_by_object_id(start_time, end_time), result)

    def test_ChaosMeshConnection_query(self):
        right_sql = f"SELECT * FROM events WHERE created_at >= '{datetime(2024, 5, 19, 8, 20, 0)}' AND created_at < '{datetime(2024, 5, 19, 8, 21, 0)}' AND reason='Applied'"
        error_sql = f"SELECT * FROM events WHERE"
        
        result = [{'id': 442, 'object_id': '929f3a12-bc6c-425a-a2ba-524ed6f2b959', 'created_at': datetime(2024, 5, 19, 8, 20), 'namespace': 'default', 'name': 'pod-failure-1-nj2tv', 'kind': 'PodChaos', 'type': 'Normal', 'reason': 'Applied', 'message': 'Successfully apply chaos for default/checkoutservice-6846846596-tff57'}]
        # 正确的 sql 语句
        self.assertEqual(cm_conn.query(right_sql), result)
        # 错误的 sql 语句
        with self.assertRaises(Exception):
            cm_conn.query(error_sql)
        
    def test_ChaosMeshConnection_create_ground_truth(self):
        start_time = datetime(2024, 5, 19, 8, 20, 0)
        end_time = datetime(2024, 5, 19, 8, 21, 0)
        with open("./api/tests/choaos_mesh.json", "r") as f:
            data = json.load(f)
        self.assertEqual(cm_conn.create_ground_truth(start_time.timestamp(), end_time.timestamp()), data["create_ground_truth_result"])
        
        # with open("./api/tests/test_feature_extract_data.json", "w") as f:
        #     json.dump(data, f, indent=4)
        